import os

from utils.file_util import *
from constants import TRAIN_PATH, TEST_PATH, TOTAL_MARK_PATH, TIME_MARK_PATH
from data_process.path_algorithm.PathAlgorithm import PathAlgorithm


class SingleDirSameAlgorithm(PathAlgorithm):
    def __init__(self, params):
        self.params = params
        self.rootPath = params.rootPath
        self.volun = params.volun
        self.innerDirName = params.innerDirName
        self.vocabPath = params.vocabPath
        self.mapPath = params.mapPath
        self.trainPath = TRAIN_PATH
        self.testPath = TEST_PATH
        self.selectTrainNum = params.selectTrainNum

    def generateTimeMarkFile(self):
        markDir = os.path.join(self.rootPath, str(self.volun), self.innerDirName)
        markPath = os.path.join(markDir, TIME_MARK_PATH)
        with open(markPath, 'rb') as f:
            markData = pickle.load(f)
        timeMarkMap = {}
        for key in markData.keys():
            savePath = os.path.join(markDir, key)
            timeMarkMap[savePath] = markData[key]
        with open(TOTAL_MARK_PATH, 'wb') as f:
            pickle.dump(timeMarkMap, f)

    def generatePath(self):
        # createPathCorpus = os.path.join(self.rootPath, str(self.volun), self.innerDirName)
        createPathCorpus = self.innerDirName
        mapTabel = CreateMapTable(self.mapPath)
        SplitDataPath(createPathCorpus, self.vocabPath, 0.2)
        self.selectDataFromFile(self.trainPath, self.selectTrainNum)
        self.addAllPrefix()
        createContentFromFilePath(self.trainPath, self.testPath, mapTabel)
        # self.generateTimeMarkFile()

    def selectDataFromFile(self, filePath, selectNumber=-1):
        allPath = readDataFromFile(filePath)
        prefixDict = {}
        trainPath = []
        for curPath in allPath:
            pre = self.getPrefixFromPath(curPath)
            if pre not in prefixDict.keys():
                prefixDict[pre] = []
            prefixDict[pre].append(curPath)
        for curPre in prefixDict.keys():
            curList = prefixDict[curPre]
            curTrain = []
            if selectNumber>0:
                fNum = min(len(curList), selectNumber)
                offList = random.sample(range(0, len(curList)), fNum)
                for i in offList:
                    curTrain.append(curList[i])
            else:
                curTrain = curList
            trainPath += curTrain

        with open(self.trainPath, 'w') as f:
            for path in trainPath:
                f.write(f'{path}\n')


    def addAllPrefix(self):
        self.addPrefixOfPathFile(self.trainPath)
        self.addPrefixOfPathFile(self.testPath)

    def addPrefixOfPathFile(self, filePath):
        savePaths = []
        # prefix = os.path.join(self.rootPath, str(self.volun), self.innerDirName)
        prefix = self.innerDirName
        paths = readDataFromFile(filePath)
        for curPath in paths:
            res = os.path.join(prefix, curPath)
            savePaths.append(res)
        with open(filePath, 'w') as f:
            for curPath in savePaths:
                f.write(f'{curPath}\n')

    def getPrefixFromPath(self, path):
        pathLen = len(path)
        prefix = ''
        for i in range(pathLen-1, 0, -1):
            if(path[i] == '_'):
                prefix = path[0:i]
                break
        return prefix

